module Zeno.Solver
(
  solver
)
where

import Prelude ()
import StdImports
import qualified Data.Set as Set

import Zeno.Core
import Zeno.Type
import Zeno.Term
import Zeno.Clause
import Zeno.Function
import Zeno.Evaluation
import Zeno.Checker
import Zeno.ProofLogic ( Proof (..) )
import qualified Zeno.ProofLogic as PL

{- 
TODO

Command line option for randomSequence vs sequence

Figure out why x <= y + x isn't solving properly i.e., by using hcn

Same term, different path critical term finding is merged into one, 
  so it doesn't get done twice

Do things in phases?
  like after induction there is no generalisation?
  after splitting there is no induction?
  
Remember to add kterm hypotheses to the conditions
  for elem-min
  
Theory about removal of multiple linked induction hypotheses on application

Theory about only adding a hypothesis one the "first" inductive step of a chain
  that is to say do not add a hypothesis if hypotheses already exist
  
Try merging case-splitting back into induction
  
just remember to apply function definitions whilst finding critical terms
  
Try iding each case-analysis and never use case-analysis to do induction on
  a child of a variable that was created with this analysis
  
Try "hypothesis can be proven from conditions of goal using the solver"
  rather than it can directly be inferred

Should variable sources distinguish between case-split and induction generated ones?
  i.e. data VariableSource = CaseSplit [Id] | Induction [Id] 
  
-}

data SolverState
  = SolverState   { sZeno :: !Zeno,
                    sGoal :: !ZClause,
                    sHypotheses :: ![ZQuantifiedClause],
                    sRandom :: !StdGen }
    
type Solver = Reader SolverState
type StatefulSolver = State SolverState
type SolverStep = Solver [ZProof]

instance Contains SolverState Zeno where
  retrieve = sZeno
  update z s = s { sZeno = z } 
  
instance IdCounter SolverState where
  newId ss = 
    let zeno = sZeno ss
        (id, zeno') = newId zeno
        ss' = ss { sZeno = zeno' }
    in (id, ss')

instance WithinTraversable ZTerm SolverState where
  mapWithinM f s = do
    goal' <- mapWithinM f (sGoal s)
    hyps' <- mapM (mapWithinM f) (sHypotheses s)
    return $ s { sGoal = goal', sHypotheses = hyps' }
  
solver :: String -> StdGen -> ZenoReader ZProof
solver lemma_name random = do
  func <- lookupFunction lemma_name
  let lemmas = clausifyFunction func
  when (null lemmas) $ error $ "Lemma " ++ lemma_name ++ " not found."    
  false <- anyM falsifiable lemmas
  when false $ error "Lemma is false, counter-example found."
  proofs <- mapM solver' lemmas
  return $ case proofs of
    [] -> Unknown
    [proof] -> proof
    proofs -> PL.conjoinStep (PL.ProofStep "proofs" Nothing []) proofs
  where
    solver' :: ZClause -> ZenoReader ZProof
    solver' lemma = do
      zeno <- ask
      let state = traceZ ("solving " ++ show lemma) $
            SolverState { sZeno = zeno,
                          sGoal = lemma,
                          sHypotheses = mempty,
                          sRandom = random }
      return
        $ flip runReader state 
        $ addStep "goal" [] 
        $ refreshGoalSources 
        $ solve_def solve_start

combine :: [Solver ZProof] -> Solver ZProof
combine = fmap PL.anyProven . sequence --randomSequence 

solve_start :: Solver ZProof
solve_start = solve_def 
        --    $ traceGoal "goal"
            $ solve_eql
            $ solve_con
            $ combine [ solve_fac
                      , solve_ind
                      , solve_cse
                      , solve_hgn
                      , solve_hyp
                      , solve_gen 
                      ]

solve_def :: Solver ZProof -> Solver ZProof
solve_def cont = do
  Clause eq conds <- asks sGoal
  hyps <- asks sHypotheses
  let eq' = substituteWithin (equalitiesToSubstitution conds) eq 
      (goal', any_changed) = runWriter
        $ applyTermDefinitions conds (Clause eq' conds)
      hyps' = id
        $ filter (not . equal . clauseEquality . quantifiedObject)
        $ fst 
        $ runWriter 
        $ mapM (applyTermDefinitions conds) hyps
  local (\s -> s { sHypotheses = hyps' }) 
    $ if not (getAny any_changed)
        then cont
        else localGoal goal'
                $ addStep "def" []
                $ cont

solve_eql :: Solver ZProof -> Solver ZProof
solve_eql cont = do
  Clause eq@(t1 `Equals` t2) conds <- asks sGoal
  if equalTerms t1 t2
      then return $ Proven (PL.ProofStep "eql" Nothing []) []
    else if unequalTerms t1 t2
      then solve_con (return Unknown)
    else do
      let conds' = filter (uncurry (/=) . equalityToPair) conds
      localGoal (Clause (t1 `Equals` t2) conds') cont
   
solve_con :: Solver ZProof -> Solver ZProof
solve_con cont = do
  conds <- clauseConditions <$> asks sGoal
  if any unequal conds
    then return $ Proven (PL.ProofStep "con" Nothing []) []
    else cont
 
generalise :: ZTerm -> Solver ZProof
generalise term = do
  stateful (extendState newIdS) $ \new_id -> do
    goal <- asks sGoal
    let new_var = ZVar new_id Nothing (termType term) (ZVariable [])
        step = "gen " ++ show term ++ " => " ++ show new_var
    local (substituteTakingSources term (Var new_var))
      $ checkFalsifiable
      $ addStep step []
      $ removeAllHypotheses
      $ solve_start
      
solve_hyp :: Solver ZProof
solve_hyp = PL.anyProven <$> applyHypotheses maybeApplyHyp
  where
    maybeApplyHyp :: ZQuantifiedClause -> Solver [ZProof]
    maybeApplyHyp hyp = do
      goal@(l_goal `Equals` r_goal) <- clauseEquality <$> asks sGoal
      let flipped_hyp = flipQuantifiedClauseEquality hyp
          apply_left = findApplications l_goal hyp
          apply_right = findApplications r_goal flipped_hyp
          l_eqs = map (\t -> (t `Equals` r_goal)) apply_left
          r_eqs = map (\t -> (l_goal `Equals` t)) apply_right
          all_eqs = l_eqs ++ r_eqs
      let step = "hyp " ++ show hyp
          solveEq eq = localGoalEquality eq
                     $ checkFalsifiable
                     $ addStep step []
                     $ solve_start
      mapM solveEq all_eqs
    
    findApplications :: ZTerm -> ZQuantifiedClause -> [ZTerm]
    findApplications goal_term
                     (Quantified vars (Clause (from_hyp `Equals` to_hyp) [])) = 
      map applyUnifier $ 
        filter (completeUnifier vars) $ 
          allUnifiers from_hyp goal_term
      where 
        applyUnifier :: ZTermSubstitution -> ZTerm
        applyUnifier uni = 
          let from_hyp' = substituteWithin uni from_hyp 
              to_hyp' = substituteWithin uni to_hyp
          in substituteTakingSources from_hyp' to_hyp' goal_term
          
          
solve_fac :: Solver ZProof
solve_fac = do
  goal <- sGoal <$> ask
  let (l_term `Equals`  r_term) = clauseEquality goal
      (l_func : l_args) = flattenTerm l_term
      (r_func : r_args) = flattenTerm r_term
  if l_func /= r_func
    then return Unknown
    else do
      step <- proofstep "fct" []
      let factor_goals = map pairToEquality 
            $ l_args `zip` r_args
          factor_proofs = flip mapM factor_goals $
            \eq -> localGoalEquality eq 
                      $ checkFalsifiable
                      $ addStep "fac" [] 
                      $ solve_start
      PL.allProvenStep step <$> factor_proofs  
      
solve_hgn :: Solver ZProof
solve_hgn = PL.anyProven <$> applyHypotheses generaliseHyp
  where
    generaliseHyp :: ZQuantifiedClause -> Solver [ZProof]
    generaliseHyp (Quantified vars (Clause eq@(l_hyp `Equals` r_hyp) [])) 
      | not (constructorTerm r_hyp) = return [Unknown]
      | otherwise = do
          goal_term <- equalityLeft <$> clauseEquality <$> sGoal <$> ask 
          
          let possibleTerm :: ZTerm -> Bool
              possibleTerm term 
                 = not (isVar term) 
                && functionTerm term 
                && destructibleTerm term 
                && length (variables term) == length vars + 1
    
          sequence $ do
            possible <- filter possibleTerm (strictlyWithinList l_hyp)
            uni <- allUnifiers possible goal_term
            if not (completeUnifier vars uni) 
              then return (return Unknown)
              else do
                let eq' = substituteWithin uni eq
                    genr = substituteWithin uni possible
                return $ addCond eq' 
                       $ addStep "hgn" []
                       $ generalise genr

solve_gen :: Solver ZProof
solve_gen = do
  inds <- filter variableTerm
      <$> map fst
      <$> filter (not . usableCriticalTerm)
      <$> criticalTerms
  Clause (l_eq `Equals` r_eq) conds <- asks sGoal
  let -- cons = map Var (constructors goal)
      subterms = concatMap strictlyWithinList (l_eq : r_eq : map equalityLeft conds)
      superterms = nub $
        concatMap (usableSuperterms subterms) inds 
        -- (inds ++ cons)
  PL.anyProven <$> mapM generalise superterms
  where  
    usableSuperterms :: [ZTerm] -> ZTerm -> [ZTerm]
    usableSuperterms all_terms sub_term = 
      removeSuperterms (filter superterm all_terms)
      --removeSubterms all_duplicates
      where
        all_duplicates :: [ZTerm]
        all_duplicates = filter superterm (duplicates all_terms)

        superterm :: ZTerm -> Bool
        superterm t = destructibleTerm t && isApp t 
          && t `contains` sub_term && functionTerm t

solve_cse :: Solver ZProof
solve_cse = do
  crits <- filter usableCriticalTerm
    <$> mergeCriticalTerms
    <$> criticalTerms
  case find (isApp . fst) crits of
    Nothing -> return Unknown
    Just split -> goSplit split
  where
    provenPath :: ZTerm -> ZTerm -> Solver Bool
    provenPath term kterm =
      if isApp kterm
        then return False
        else do
          conds <- clauseConditions <$> asks sGoal
          let goal = Clause (term `Equals` kterm) conds
          proof <- localGoal goal
                $ checkFalsifiable
                $ removeAllHypotheses
                $ refreshGoalSources
                $ solve_start
          return (PL.proven proof)

    goSplit :: CriticalTerm -> Solver ZProof
    goSplit cterm@(term, source) = do
      case instantiateConstructors cterm of
        Nothing -> return Unknown
        Just cons_s -> stateful cons_s $ \cons -> do
          let kterms = map fst cons
              sourced_term = addSources [source] <$> term
          maybe_proven <- findM (provenPath sourced_term) kterms
          case maybe_proven of
            Just kterm -> do
              step <- proofstep "pcs" []
              goSplit' kterm
            Nothing -> do
              splits <- mapM goSplit' kterms
              step <- proofstep "cse" []
              return $ PL.allProvenStep step splits
      where
        goSplit' :: ZTerm -> Solver ZProof
        goSplit' kterm = do
          Clause eq conds <- sGoal <$> ask
          let goal = Clause eq ((term `Equals` kterm) : conds)
              step = "cse " ++ show term ++ " => " ++ show kterm
          localGoal goal
            $ addStep step []
            $ solve_start

solve_ind :: Solver ZProof
solve_ind = do
  crits <- filter usableCriticalTerm 
    <$> mergeCriticalTerms
    <$> criticalTerms 
  case find (variableTerm . fst) crits of
    Nothing -> return Unknown
    Just ind -> goInductive ind
  where
    goInductive :: CriticalTerm -> Solver ZProof
    goInductive cterm@(ind_term@(Var ind_var), _) = filterConditions $ do
      case instantiateConstructors cterm of
        Nothing -> return Unknown
        Just cons_s -> stateful cons_s $ \cons -> do
          ind_proofs <- mapM (goInductive' ind_var) cons
          step <- proofstep "ind" []
          return (PL.allProvenStep step ind_proofs) 
    
    filterConditions :: Solver a -> Solver a
    filterConditions solver = do
      Clause eq conds <- asks sGoal
      let vars = map variables (eq : conds)
          singleton var = (<= 1)
                        $ length
                        $ filter (elem var) vars
          usable cond = not
                      $ any singleton
                      $ variables cond
          conds' = filter usable conds
      localGoal (Clause eq conds') solver
  
    purifyGoalConditions :: Solver a -> Solver a
    purifyGoalConditions solver = do
      goal <- asks sGoal
      purified <- extendReader (purifyConditions goal)
      localGoal purified solver
    
    conditionPowerset :: Solver a -> Solver [a]
    conditionPowerset solver = do
      Clause eq ants <- asks sGoal
      let goals = map (Clause eq) (powerset ants)
      mapM (flip localGoal solver) goals
      
    goInductive' :: ZVar -> (ZTerm, [ZVar]) -> Solver ZProof 
    goInductive' ind_var (new_term, ind_vars) = do
      goal <- asks sGoal
      hyps <- asks sHypotheses
      let others = filter (/= ind_var) (variables goal)
          freshen = filter (not . isFunctionType . zVarType) others
          
          createHyp :: ZVar -> StatefulSolver ZQuantifiedClause
          createHyp rec_var = extendState $ do
            --fresh <- mapM freshVariable freshen
            let sublist = (ind_var, rec_var) : [] -- : (freshen `zip` fresh)
                hyp = substitute (Set.fromList sublist) goal
            return (Quantified freshen hyp)
           
          groundHypothesis :: ZQuantifiedClause -> [ZQuantifiedClause]
          groundHypothesis qcls@(Quantified vars cls) =
            if not (ind_var `elem` vars)
              then [qcls]
              else let new_vars = filter (/= ind_var) vars
                       generateHyp new_term = Quantified new_vars 
                         $ substituteTakingSources (Var ind_var) new_term cls
                   in map generateHyp (new_term : map Var ind_vars)
               {-      if null ind_vars 
                       then [generateHyp new_term]
                       else map (generateHyp . Var) ind_vars  -}

      let new_goal = substituteTakingSources (Var ind_var) new_term goal
          new_hyps = concatMap groundHypothesis hyps
          step = "ind " ++ show ind_var ++ " => " ++ show new_term
                
      stateful (mapM createHyp ind_vars) $ \created_hyps -> id
        $ local (\s -> s { sGoal = new_goal })
        $ local (\s -> s { sHypotheses = created_hyps ++ new_hyps })
        $ addStep step created_hyps 
        $ solve_start
       
traceGoal' tag solver = solver
        
traceGoal :: String -> Solver ZProof -> Solver ZProof
traceGoal tag solver = do
  goal <- asks sGoal
  hyps <- asks sHypotheses
  let hyps_s = intercalate "\n" $ map show hyps
  proof <- solver
  return (trace (tag ++ ":" ++ show goal ++ " with\n" ++ hyps_s ++ " is " ++ show (PL.proven proof) ++ "\n") proof)

checkFalsifiable :: Solver ZProof -> Solver ZProof
checkFalsifiable solver = do
  goal <- asks sGoal
  false <- extendReader (falsifiable (removeSources <$> goal))
  if false then return Unknown else solver

-- Do we need to refresh the variables in the hypotheses?
refreshGoalSources :: Solver a -> Solver a
refreshGoalSources solver = do
  goal <- asks sGoal
  localGoal (removeSources <$> goal) solver
  
completeUnifier :: [ZVar] -> ZTermSubstitution -> Bool
completeUnifier vars uni =
  Set.fromList (map Var vars) == Set.mapMonotonic fst uni

criticalTerms :: Solver [CriticalTerm]
criticalTerms = do
  Clause (l_term `Equals` r_term) conds <- sGoal <$> ask
  let cond_lefts = map equalityLeft conds
      all_terms = l_term : r_term : cond_lefts
      
      -- This is to stop infinite loops in provenPath case splitting
      termExists (term, _) = isVar term
        || not (any (flip contains term) all_terms)
  return 
    $ filter termExists
    $ nub
    $ filter (not . flip elem cond_lefts . fst)
    $ filter (destructibleTerm . fst)
    $ mapMaybe (runWriterT . criticalTerm)
    $ all_terms
  where
    criticalTerm :: ZTerm -> WriterT [Id] Maybe ZTerm
    criticalTerm term =
      case applyFunction [] term of
        Nothing -> lift (Just term)
        Just (Value {}) -> lift Nothing
        Just (CaseOf id strict_term _) -> do
          tell [id]
          if not (term `contains` strict_term)
            then lift (Just strict_term)
            else criticalTerm strict_term
             

-- Applicable hypotheses are those with all conditions fulfilled
--   the given hypotheses should be conditionless
applyHypotheses :: (ZQuantifiedClause -> Solver [ZProof]) -> Solver [ZProof]
applyHypotheses apply = do
  hyps <- asks sHypotheses
  stateful (mapM freshenHypothesis hyps) 
    $ \fresh_hyps -> concat <$> concatMapM applyHyp (hyps `zip` fresh_hyps)
  where
    freshenHypothesis :: ZQuantifiedClause -> StatefulSolver ZQuantifiedClause
    freshenHypothesis (Quantified vars cls) = do
      fresh_vars <- extendState $ mapM freshVariable vars
      let sub = vars `zip` fresh_vars
          fresh_cls = fmap (substituteVars sub) cls
      return (Quantified fresh_vars fresh_cls)
      where
        substituteVars :: [(ZVar, ZVar)] -> ZVar -> ZVar
        substituteVars vars var = 
          case lookup var vars of
            Nothing -> var
            Just new_var -> addSources (variableSources var) new_var
  
    applyHyp :: (ZQuantifiedClause, ZQuantifiedClause) -> Solver [[ZProof]]
    applyHyp (old_cls, fresh_cls) = do
      apps <- applicable fresh_cls
      removeHypothesis old_cls (mapM apply apps)
  
    applicable :: ZQuantifiedClause -> Solver [ZQuantifiedClause]
    applicable cls@(Quantified _ (Clause _ [])) = return [cls]
    applicable (Quantified c_vars (Clause c_eq c_conds)) = do
      unis <- mapM removeCondition c_conds
      let unis' = map Set.unions (sequence unis)
      return $ do
        uni <- unis'
        guard (isOneToMany uni)
        let vars = map (fromVar . fst) (toList uni)
            c_vars' = filter (not . flip elem vars) c_vars
            c_eq' = substituteWithin uni c_eq
        return 
          $ Quantified c_vars'
          $ Clause c_eq' []
      where
        removeCondition :: ZEquality -> Solver [ZTermSubstitution]
        removeCondition eq@(c_l `Equals` c_r) = do
          g_conds <- clauseConditions <$> asks sGoal
          return $ do
            g_l `Equals` g_r <- g_conds
            let uni_l = unify c_l g_l
                uni_r = unify c_r g_r
                uni = uni_l `Set.union` uni_r
            if isOneToMany uni && validUnifier uni
              then [uni]
              else []
          where 
            validUnifier :: ZTermSubstitution -> Bool
            validUnifier = all (flip elem (map Var c_vars) . fst)

addStep :: String -> [ZQuantifiedClause] -> Solver ZProof -> Solver ZProof
addStep name added solver = do
  goal <- sGoal <$> ask
  PL.addStep (PL.ProofStep name (Just goal) added) <$> solver
  
  
invalidTerm :: ZTerm -> Bool
invalidTerm term = 
  let ((Var func) : args) = flattenTerm term
      arg_types = flattenFunType (zVarType func)
      res_types = drop (length args) arg_types 
  in null res_types
       
randomSequence :: [Solver a] -> Solver [a]
randomSequence solvers = do
  random <- sRandom <$> ask
  let (solvers', random') = shuffle solvers random
  localRandom random' (splitGens solvers')
  where
    localRandom :: StdGen -> Solver a -> Solver a
    localRandom gen = local (\s -> s { sRandom = gen })
        
    splitGens :: [Solver a] -> Solver [a]
    splitGens = foldl' splitGen (return [])
      where 
        splitGen :: Solver [a] -> Solver a -> Solver [a]
        splitGen ss s = do
          random <- sRandom <$> ask
          let (r1, r2) = split random
          ss' <- localRandom r2 ss
          s' <- localRandom r1 s
          return (s' : ss')

localGoal :: ZClause -> Solver a -> Solver a
localGoal goal = local $ \s -> s { sGoal = goal }

localGoalEquality :: ZEquality -> Solver a -> Solver a
localGoalEquality eq = local $ \s -> s 
  { sGoal = (sGoal s) { clauseEquality = eq } }

localGoalConditions :: [ZEquality] -> Solver a -> Solver a
localGoalConditions conds = local $ \s -> s
  { sGoal = (sGoal s) { clauseConditions = conds } }
  
addCond :: ZEquality -> Solver a -> Solver a
addCond cond = local $ \s -> s
  { sGoal = (sGoal s) 
    { clauseConditions = cond : (clauseConditions (sGoal s)) } }

addConditions :: [ZEquality] -> Solver a -> Solver a
addConditions = foldl (\f c -> f . addCond c) id
    
addHypothesis :: ZQuantifiedClause -> Solver a -> Solver a
addHypothesis cls = local $ \s -> s
  { sHypotheses = cls : (sHypotheses s) }
  
removeHypothesis :: ZQuantifiedClause -> Solver a -> Solver a
removeHypothesis cls = local $ \s -> s
  { sHypotheses = delete cls (sHypotheses s) }
  
removeHypotheses :: [ZQuantifiedClause] -> Solver a -> Solver a
removeHypotheses = foldl (\f c -> f . removeHypothesis c) id

removeAllHypotheses :: Solver a -> Solver a
removeAllHypotheses = local $ \s -> s { sHypotheses = [] }

addHypotheses :: [ZQuantifiedClause] -> Solver a -> Solver a
addHypotheses = foldl (\f c -> f . addHypothesis c) id
       
pureRecordType :: ZType -> Bool
pureRecordType z_type = 
  let ((TypeVar dtype@(ZDataType {})) : _) = flattenAppType z_type
  in  not (recursiveTypeVar dtype)
  
proofstep :: String -> [ZQuantifiedClause] -> Solver ZProofStep
proofstep name added = do
  goal <- sGoal <$> ask
  return (PL.ProofStep name (Just goal) added) 
